/*
 *   Copyright 2021 zzh
 *
 *   Licensed under the Apache License, Version 2.0 (the "License");
 *   you may not use this file except in compliance with the License.
 *   You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 *   Unless required by applicable law or agreed to in writing, software
 *   distributed under the License is distributed on an "AS IS" BASIS,
 *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *   See the License for the specific language governing permissions and
 *   limitations under the License.
 *
 */

package com.aduib.boot.common.asm;

import com.aduib.boot.common.concurrent.JctoolsHelper;
import com.aduib.boot.common.log.LOG;
import com.aduib.boot.common.util.CharUtils;
import com.aduib.boot.common.util.ClassUtils;
import com.aduib.boot.common.util.FileUtils;
import com.aduib.boot.common.util.IOUtils;
import com.aduib.boot.common.util.StrUtils;
import com.aduib.boot.common.util.SystemUtils;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.Type;

import java.io.ByteArrayInputStream;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Objects;

import static org.objectweb.asm.Opcodes.ALOAD;
import static org.objectweb.asm.Opcodes.ARETURN;
import static org.objectweb.asm.Opcodes.INVOKESPECIAL;
import static org.objectweb.asm.Opcodes.RETURN;

/**
 * @description: AsmClassGenerator
 * @author: zzh
 * @date: 2021/10/18 10:29
 */
public class AsmClassGenerator {

  private final boolean enabledPrintClassInfo =
      Boolean.parseBoolean(SystemUtils.getStringValue("aduib.class.enabledPrint", "false"));

  private final AsmClassGeneratorClassLoader generatorClassLoader =
      new AsmClassGeneratorClassLoader();
  private final List<Class<?>> interfaces = new ArrayList<>();
  private final List<Class<?>> genericTypes = new ArrayList<>();
  private final List<ConstructorDescriptor> constructorDescriptorList = new ArrayList<>();
  private final List<FieldDescriptor> fieldDescriptorList = new ArrayList<>();
  private final List<MethodDescriptor> methodDescriptorList = new ArrayList<>();
  private String className;
  private int modifier = Opcodes.ACC_PUBLIC;
  private int typeValue = Opcodes.ACC_SUPER;
  private Class<?> superClass = Object.class;

  public void addInterface(Class<?> interfaceClass) {
    this.interfaces.add(interfaceClass);
  }

  public String getClassName() {
    return className;
  }

  public void setClassName(String className) {
    this.className = className;
  }

  public Class<?> getSuperClass() {
    return superClass;
  }

  public void setSuperClass(Class<?> superClass) {
    this.superClass = superClass;
  }

  public List<Class<?>> getInterfaces() {
    return interfaces;
  }

  public int getModifier() {
    return modifier;
  }

  public void setModifier(List<Opcode> opcodes) {
    this.modifier = opcodes.stream().map(Opcode::getOpcodes).reduce(0, Integer::sum);
  }

  public int getTypeValue() {
    return typeValue;
  }

  public void setTypeValue(Opcode opcode) {
    this.typeValue = opcode.getOpcodes();
  }

  /**
   * 获取类对应内部名
   *
   * @param clazz
   * @return
   */
  private String getClassInternalName(Class<?> clazz) {
    if (Objects.isNull(clazz)) {
      return null;
    }
    return Type.getInternalName(clazz);
  }

  /**
   * 获取类对应的字节码描述信息
   *
   * @param clazz
   * @return
   */
  private String getClassDescriptor(Class<?> clazz) {
    if (Objects.isNull(clazz)) {
      return null;
    }
    return Type.getDescriptor(clazz);
  }

  /**
   * 获取泛型对应的字节码描述信息
   *
   * @param genericTypes
   * @return
   */
  private String getGenericTypesDescriptor(List<Class<?>> genericTypes) {
    if (genericTypes.isEmpty()) {
      return null;
    }
    StringBuilder sb = new StringBuilder();
    sb.append("<");
    for (Class<?> genericType : genericTypes) {
      sb.append(Type.getDescriptor(genericType));
    }
    sb.append(">");
    return sb.toString();
  }

  /**
   * 批量生成接口的字节码描述内容
   *
   * @param interfaceTypes
   * @return
   */
  private String[] getInterfaceTypesDescriptor(List<Class<?>> interfaceTypes) {
    if (interfaceTypes.isEmpty()) {
      return null;
    }
    String[] interfaceDescs = new String[interfaceTypes.size()];
    for (int i = 0; i < interfaceTypes.size(); i++) {
      interfaceDescs[i] = Type.getInternalName(interfaceTypes.get(i));
    }
    return interfaceDescs;
  }

  /**
   * 返回加载的class
   *
   * @return
   */
  public Class<?> build() throws IOException, ClassNotFoundException {
    ClassWriter classWriter = new ClassWriter(ClassWriter.COMPUTE_MAXS);
    AsmClassGeneratorVisitor visitor = new AsmClassGeneratorVisitor(classWriter);
    // 生成类声明信息
    String internalClassName = this.className.replaceAll("\\" + StrUtils.DOT, StrUtils.SLASH);
    visitor.visit(
        Opcodes.V1_8,
        modifier | typeValue,
        internalClassName,
        getGenericTypesDescriptor(genericTypes),
        getClassInternalName(superClass),
        getInterfaceTypesDescriptor(interfaces));
    // 生成构造器
    for (ConstructorDescriptor constructorDescriptor : constructorDescriptorList) {
      MethodVisitor methodVisitor =
          visitor.visitMethod(
              constructorDescriptor.getModifier(),
              constructorDescriptor.getMethodName(),
              Type.getMethodDescriptor(Type.VOID_TYPE, constructorDescriptor.getParameterTypes()),
              null,
              constructorDescriptor.getThrowables());
      constructorDescriptor.visitCode(methodVisitor);
    }
    // 生成属性
    for (FieldDescriptor fieldDescriptor : fieldDescriptorList) {
      visitor.visitField(
          fieldDescriptor.getModifier(),
          fieldDescriptor.getFieldName(),
          getClassDescriptor(fieldDescriptor.getFieldType()),
          getGenericTypesDescriptor(fieldDescriptor.getGenericTypes()),
          fieldDescriptor.getValue());
    }
    // 生成方法
    for (MethodDescriptor methodDescriptor : methodDescriptorList) {
      MethodVisitor methodVisitor =
          visitor.visitMethod(
              methodDescriptor.getModifier(),
              methodDescriptor.getMethodName(),
              Type.getMethodDescriptor(
                  methodDescriptor.getReturnType(), methodDescriptor.getParameterTypes()),
              getGenericTypesDescriptor(methodDescriptor.getGenericTypes()),
              methodDescriptor.getThrowables());
      methodDescriptor.visitCode(methodVisitor);
    }
    visitor.visitEnd();
    byte[] javaFile = classWriter.toByteArray();
    generatorClassLoader.add(className, javaFile);
    if (enabledPrintClassInfo) {
      FileUtils.writeFromStream(
          new ByteArrayInputStream(javaFile),
          SystemUtils.INSTANCE.getHomeDir() + File.separator + className + FileUtils.EXT_CLASS);
    }
    Class<?> clazz = ClassUtils.forName(className, generatorClassLoader);
    LOG.d("AsmClassGenerator{1}:" + clazz);
    LOG.d("AsmClassGenerator{2}:" + className);
    return clazz;
  }

  public void addConstructorDescriptor(ConstructorDescriptor constructorDescriptor) {
    constructorDescriptorList.add(constructorDescriptor);
  }

  public void addFieldDescriptor(FieldDescriptor fieldDescriptor) {
    fieldDescriptorList.add(fieldDescriptor);
  }

  public void addMethodDescriptor(MethodDescriptor methodDescriptor) {
    methodDescriptorList.add(methodDescriptor);
  }

  public enum Opcode {
    OC_PUBLIC(Opcodes.ACC_PUBLIC),
    OC_PRIVATE(Opcodes.ACC_PRIVATE),
    OC_PROTECTED(Opcodes.ACC_PROTECTED),
    OC_STATIC(Opcodes.ACC_STATIC),
    OC_FINAL(Opcodes.ACC_FINAL),
    OC_CLASS(Opcodes.ACC_SUPER);

    private final int opcodes;

    Opcode(int opcodes) {
      this.opcodes = opcodes;
    }

    public int getOpcodes() {
      return opcodes;
    }
  }

  /** 自定义方法内容创建 */
  interface CustomizeMethodContentGenerator {

    boolean beforeCustomize(MethodVisitor methodVisitor);

    boolean generate(MethodVisitor methodVisitor, boolean enabledCustomize);

    boolean afterCustomize(MethodVisitor methodVisitor, boolean enabledReturn);

    default void visitCode(MethodVisitor methodVisitor) {
      afterCustomize(methodVisitor, generate(methodVisitor, beforeCustomize(methodVisitor)));
    }
  }

  public abstract static class ConstructorDescriptor implements CustomizeMethodContentGenerator {

    private final String methodName = "<init>";

    private final int modifier;

    private final List<Class<?>> parameterTypes = new ArrayList<>();

    private final List<Class<?>> throwables = new ArrayList<>();

    public ConstructorDescriptor(List<Opcode> modifiers) {
      this.modifier = modifiers.stream().map(Opcode::getOpcodes).reduce(0, Integer::sum);
    }

    public ConstructorDescriptor addParameterType(Class<?> parameterType) {
      this.parameterTypes.add(parameterType);
      return this;
    }

    public ConstructorDescriptor addThrowable(Class<?> throwable) {
      this.throwables.add(throwable);
      return this;
    }

    public int getModifier() {
      return modifier;
    }

    public Type[] getParameterTypes() {
      Type[] types = new Type[parameterTypes.size()];
      for (int i = 0; i < parameterTypes.size(); i++) {
        types[i] = Type.getType(parameterTypes.get(i));
      }
      return types;
    }

    public String[] getThrowables() {
      String[] throwable = new String[throwables.size()];
      for (int i = 0; i < throwables.size(); i++) {
        throwable[i] = Type.getInternalName(throwables.get(i));
      }
      return throwable;
    }

    public String getMethodName() {
      return methodName;
    }

    @Override
    public boolean beforeCustomize(MethodVisitor mv) {
      mv.visitVarInsn(ALOAD, 0);
      mv.visitMethodInsn(INVOKESPECIAL, "java/lang/Object", "<init>", "()V", false);
      return true;
    }

    @Override
    public abstract boolean generate(MethodVisitor mv, boolean enabledCustomize);

    @Override
    public boolean afterCustomize(MethodVisitor mv, boolean enabledReturn) {
      if (enabledReturn) {
        mv.visitInsn(RETURN);
      }
      mv.visitMaxs(2, parameterTypes.size() + 1);
      mv.visitEnd();
      return false;
    }
  }

  public abstract static class MethodDescriptor extends ConstructorDescriptor {

    private final String methodName;

    private final Class<?> returnType;

    private final List<Class<?>> genericTypes = new ArrayList<>();

    public MethodDescriptor(List<Opcode> modifiers, String methodName, Class<?> returnType) {
      super(modifiers);
      this.methodName = methodName;
      this.returnType = returnType;
    }

    public MethodDescriptor(
        List<Opcode> modifiers,
        String methodName,
        Class<?> returnType,
        Class<?>[] parameterTypes,
        Class<?>[] exceptionTypes) {
      super(modifiers);
      this.methodName = methodName;
      this.returnType = returnType;
      if (Objects.nonNull(parameterTypes)) {
        for (Class<?> parameterType : parameterTypes) {
          addParameterType(parameterType);
        }
      }
      if (Objects.nonNull(exceptionTypes)) {
        for (Class<?> expectedType : exceptionTypes) {
          addThrowable(expectedType);
        }
      }
    }

    public MethodDescriptor addGenericType(Class<?> genericType) {
      this.genericTypes.add(genericType);
      return this;
    }

    @Override
    public String getMethodName() {
      return methodName;
    }

    public Type getReturnType() {
      return Type.getType(returnType);
    }

    public List<Class<?>> getGenericTypes() {
      return genericTypes;
    }

    @Override
    public boolean beforeCustomize(MethodVisitor mv) {
      return true;
    }

    @Override
    public abstract boolean generate(MethodVisitor mv, boolean enabledCustomize);

    @Override
    public boolean afterCustomize(MethodVisitor mv, boolean enabledReturn) {
      if (enabledReturn) {
        if (returnType.equals(void.class) || returnType.equals(Void.class)) {
          mv.visitInsn(RETURN);
        } else {
          mv.visitInsn(ARETURN);
        }
      }
      mv.visitMaxs(2, getParameterTypes().length + 1);
      mv.visitEnd();
      return false;
    }
  }

  public static class FieldDescriptor {

    private final String fieldName;

    private final Class<?> fieldType;

    private final int modifier;

    private final List<Class<?>> genericTypes = new ArrayList<>();

    private final Object value;

    public FieldDescriptor(
        String fieldName, Class<?> fieldType, List<Opcode> modifiers, Object value) {
      this.fieldName = fieldName;
      this.fieldType = fieldType;
      this.modifier = modifiers.stream().map(Opcode::getOpcodes).reduce(0, Integer::sum);
      ;
      this.value = value;
    }

    public String getFieldName() {
      return fieldName;
    }

    public Class<?> getFieldType() {
      return fieldType;
    }

    public int getModifier() {
      return modifier;
    }

    public List<Class<?>> getGenericTypes() {
      return genericTypes;
    }

    public Object getValue() {
      return value;
    }

    public FieldDescriptor addGenericType(Class<?> genericType) {
      this.genericTypes.add(genericType);
      return this;
    }
  }

  /** 类加载器 */
  static class AsmClassGeneratorClassLoader extends ClassLoader {

    private final Map<String, byte[]> classBytes = JctoolsHelper.nonBlockingHashMap();

    @Override
    protected Class<?> findClass(String qualifiedClassName) throws ClassNotFoundException {
      byte[] file = classBytes.get(qualifiedClassName);
      if (file != null) {
        return defineClass(qualifiedClassName, file, 0, file.length);
      }
      try {
        return ClassUtils.forName(qualifiedClassName, false, getClass().getClassLoader());
      } catch (ClassNotFoundException nf) {
        return super.findClass(qualifiedClassName);
      }
    }

    AsmClassGeneratorClassLoader add(final String qualifiedClassName, final byte[] javaFile) {
      classBytes.put(qualifiedClassName, javaFile);
      return this;
    }

    @Override
    protected synchronized Class<?> loadClass(final String name, final boolean resolve)
        throws ClassNotFoundException {
      return super.loadClass(name, resolve);
    }

    @Override
    public InputStream getResourceAsStream(final String name) {
      if (name.endsWith(FileUtils.EXT_CLASS)) {
        String qualifiedClassName =
            name.substring(0, name.length() - FileUtils.EXT_CLASS.length())
                .replace(File.separatorChar, CharUtils.DOT);
        byte[] file = classBytes.get(qualifiedClassName);
        if (file != null) {
          return IOUtils.toStream(file);
        }
      }
      return super.getResourceAsStream(name);
    }

    @Override
    protected void finalize() throws Throwable {
      try {
        super.finalize();
      } finally {
        classBytes.clear();
      }
    }
  }
}
